import os
import random

# 创建文件夹路径
output_dir = r'C:\Users\ASUS\Desktop\2024年新生赛出题\2024年新生赛出题(非大一组)\E题\E题_input'
os.makedirs(output_dir, exist_ok=True)

# 设置测试数据范围
N_min, N_max = 1, 1000  # 矩阵边长范围
a_min, a_max = 1, 10**5  # 矩阵元素的值范围

# 用于判断质数
def is_prime(n):
    if n <= 1:
        return False
    if n <= 3:
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    i = 5
    while i * i <= n:
        if n % i == 0 or n % (i + 2) == 0:
            return False
        i += 6
    return True

# 随机生成矩阵并确保包含至少一个质数在对角线上
def generate_matrix(n):
    matrix = []
    primes_on_diagonals = []

    for i in range(n):
        row = []
        for j in range(n):
            num = random.randint(a_min, a_max)
            row.append(num)
            # 检查是否在主对角线或副对角线位置
            if i == j or i + j == n - 1:
                if is_prime(num):
                    primes_on_diagonals.append(num)
        matrix.append(row)

    # 确保对角线上至少有一个质数
    if not primes_on_diagonals:
        i = random.randint(0, n - 1)
        if i == n - i - 1:  # 如果是主副对角线重合的中心点
            prime_value = next(p for p in range(a_min, a_max) if is_prime(p))
            matrix[i][i] = prime_value
        else:
            prime_value = next(p for p in range(a_min, a_max) if is_prime(p))
            matrix[i][i] = prime_value
            matrix[i][n - i - 1] = prime_value

    return matrix

# 生成10个测试数据文件
for idx in range(1, 11):
    n = random.randint(N_min, min(N_max, 10))  # 使用较小的n来调试, 可改为N_max
    matrix = generate_matrix(n)
    
    # 保存到文件
    file_path = os.path.join(output_dir, f'{idx}.in')
    with open(file_path, 'w') as f:
        f.write(f"{n}\n")
        for row in matrix:
            f.write(" ".join(map(str, row)) + "\n")

print("测试数据已生成并保存在指定目录中。")